# Copyright (c) Huawei Technologies Co., Ltd. 2023. All rights reserved.
import random

from typing import Union, List

from atk.case_generator.generator.generate_types import GENERATOR_REGISTRY
from atk.case_generator.generator.base_generator import CaseGenerator

from atk.configs.case_config import InputCaseConfig, CaseConfig


@GENERATOR_REGISTRY.register("ascend_generate_aclnn_group_norm_swish_grad")
class DefaultGenerator(CaseGenerator):
    def __init__(self, config):
        super().__init__(config)
        self.dy_dtype = None
        self.dy_shape = []
        self.mean_shape = []
        self.gamma_shape = []

    def after_input_config(
            self,
            index: int,
            input_case: Union[InputCaseConfig, List[InputCaseConfig]],
    ) -> Union[InputCaseConfig, List[InputCaseConfig]]:

        ## 所有输入tensor的dtype保持一致
        if index == 0:
            self.dy_dtype = input_case.dtype

        if index == 1:
            input_case.dtype = self.dy_dtype

        if index == 2:
            input_case.dtype = self.dy_dtype

        if index == 3:
            input_case.dtype = self.dy_dtype

        if index == 4:
            input_case.dtype = self.dy_dtype

        if index == 5:
            input_case.dtype = self.dy_dtype

        if index == 6:
            if type(input_case.range_values) == list:
                input_case.range_values = random.choice(input_case.range_values)
        return input_case

    def after_case_config(self, case_config: CaseConfig) -> CaseConfig:
        if 0 in case_config.inputs[0].shape: ## 空tensor
            pass
        elif case_config.inputs[0].shape == []: ## 标量tensor
            pass
        else:
            max_dy_shape = max(case_config.inputs[0].shape)
            if max_dy_shape >= 2147483649: ## 上边界用例
                case_config.inputs[6].range_values = 1
                if case_config.inputs[0].shape[1] >= 2147483649: # C==2147483649时要满足C//group<=3000
                    case_config.inputs[6].range_values = 715827883
            elif max_dy_shape == 1: ## 下边界用例
                case_config.inputs[6].range_values = 1
            else:
                if case_config.inputs[0].shape[1] > case_config.inputs[6].range_values*3000:  # 满足C//group<=3000的约束
                    case_config.inputs[0].shape[1] = case_config.inputs[6].range_values*3000
                C = case_config.inputs[0].shape[1]
                numGroups = case_config.inputs[6].range_values
                q = C // numGroups
                r = C % numGroups
                if q == 0: # C<group
                    if random.random() > 0.5:
                        case_config.inputs[0].shape[1] = case_config.inputs[6].range_values
                    else:
                        case_config.inputs[6].range_values = 1
                        case_config.inputs[0].shape[1] = random.randint(1, 3000)
                elif r != 0: # C>group且不整除
                    case_config.inputs[0].shape[1] -= r

            case_config.inputs[1].shape = [case_config.inputs[0].shape[0], case_config.inputs[6].range_values]
            case_config.inputs[2].shape = [case_config.inputs[0].shape[0], case_config.inputs[6].range_values]
            case_config.inputs[3].shape = case_config.inputs[0].shape
            case_config.inputs[4].shape[0] = case_config.inputs[0].shape[1]
            case_config.inputs[5].shape[0] = case_config.inputs[0].shape[1]
        return case_config